//
// Created by fairy on 2025/7/27.
//

#ifndef COROUTINE_HPP
#define COROUTINE_HPP

#include <coroutine>
#include <cstddef>
#include <cstdint>

namespace utils::detail_coroutine
{
    /**
 * 静态协程池实现（零开销嵌入式版）
 * @tparam BlockSize 一个块的大小
 * @tparam BlockCount 块的数量
 */
    template<size_t BlockSize, size_t BlockCount>
    class StaticCoroutinePool {
        // 内存存储区域（按最大对齐要求）
        alignas(alignof(std::max_align_t))
        static inline uint8_t storage[BlockCount * BlockSize];

        // 位图标记使用状态（每个块1位）
        static inline uint8_t bitmap[(BlockCount + 7) / 8] = {};

        // 获取位图状态
        static bool get_bit(const size_t index)
        {
            return bitmap[index / 8] & (1 << (index % 8));
        }

        // 设置位图状态
        static void set_bit(const size_t index, bool value)
        {
            auto &byte = bitmap[index / 8];
            const uint8_t mask = 1 << (index % 8);
            byte = value ? (byte | mask) : (byte & ~mask);
        }

    public:
        StaticCoroutinePool() = delete;

        StaticCoroutinePool(const StaticCoroutinePool &) = delete;

        ~StaticCoroutinePool() = delete;

        // 分配内存块
        static void *allocate(const size_t size)
        {
            // 检查请求大小是否有效
            if (size > BlockSize) {
                return nullptr;
            }

            // 线性搜索空闲块
            for (size_t i = 0; i < BlockCount; ++i) {
                if (!get_bit(i)) {
                    set_bit(i, true);
                    return &storage[i * BlockSize];
                }
            }

            return nullptr;
        }

        // 释放内存块
        static void deallocate(void *ptr)
        {
            if (ptr == nullptr) return;

            // 计算块索引
            const auto ptr_val = reinterpret_cast<uintptr_t>(ptr);
            const auto base_val = reinterpret_cast<uintptr_t>(&storage[0]);
            const auto offset = ptr_val - base_val;

            // 验证指针有效性
            if (offset % BlockSize != 0 || offset >= BlockCount * BlockSize) {
                return;
            }
            const size_t index = offset / BlockSize;
            set_bit(index, false);
        }
    };


    /**
     * 高性能调度器（支持双模式）
     * @tparam MAX_TASKS 最大任务数量
     */
    template<uint32_t(*getTick)(), std::size_t MAX_TASKS>
    class Scheduler {
    public:
        static void add(const uint32_t when, const std::coroutine_handle<> h)
        {
            if (count >= MAX_TASKS) {
                h.resume(); // 直接恢复
                return;
            }

            // 插入并保持数组有序
            int32_t pos = static_cast<int32_t>(count) - 1;
            while (pos >= 0 && when < tasks[pos].when) {
                tasks[pos + 1] = tasks[pos]; // 后移
                --pos;
            }
            tasks[pos + 1] = {when, h};
            ++count;

            // 更新最近到期时间
            if (pos + 1 == 0) {
                next_time = tasks[0].when;
            }
        }

        /**
         * 高性能poll实现（支持双模式）
         * @tparam StepMode 是否使用分步执行模式
         */
        template<bool StepMode>
        static void poll()
        {
            const auto now = getTick();
            if (count == 0 || next_time > now) return;

            if constexpr (StepMode) {
                // 分步模式优化
                if (tasks[0].when <= now) {
                    auto h = tasks[0].h;
                    // 移动数组
                    for (size_t i = 0; i < count - 1; ++i) {
                        tasks[i] = tasks[i + 1];
                    }
                    count--;

                    next_time = count > 0 ? tasks[0].when : UINT32_MAX;

                    if (!h.done()) {
                        h.resume();
                        if (h.done()) h.destroy();
                    }
                }
            } else {
                // 批量模式优化
                size_t processed = 0;
                for (size_t i = 0; i < count;) {
                    if (tasks[i].when > now) {
                        i++;
                        continue;
                    }

                    auto h = tasks[i].h;
                    // 移除当前元素
                    for (size_t j = i; j < count - 1; j++) {
                        tasks[j] = tasks[j + 1];
                    }
                    count--;

                    if (!h.done()) {
                        h.resume();
                        if (h.done()) h.destroy();
                    }
                    processed++;
                }
                next_time = count > 0 ? tasks[0].when : UINT32_MAX;
            }
        }

        // 获取下次唤醒时间（低功耗支持）
        static uint32_t get_next_timeout()
        {
            return next_time;
        }

        Scheduler() = delete;

        ~Scheduler() = delete;

    private:
        struct Task {
            uint32_t when{};
            std::coroutine_handle<> h;
        };

        static inline Task tasks[MAX_TASKS] = {};
        static inline size_t count = 0;
        static inline uint32_t next_time = UINT32_MAX;
    };

    /**
     * 协程任务
     * @tparam FrameSize 协程帧大小
     * @note 经过实测，终端程序最少需要48个字节的协程帧
     */
    template<size_t FrameSize, size_t FrameCount>
    struct AsyncTask {
        struct promise_type {
            AsyncTask get_return_object() { return {}; }
            std::suspend_never initial_suspend() { return {}; }
            std::suspend_never final_suspend() noexcept { return {}; }

            void return_void()
            {
            }

            void unhandled_exception()
            {
            }

            // 使用静态池分配协程帧
            void *operator new(size_t size)
            {
                return StaticCoroutinePool<FrameSize, FrameCount>::allocate(size);
            }

            void operator delete(void *ptr)
            {
                StaticCoroutinePool<FrameSize, FrameCount>::deallocate(ptr);
            }
        };
    };


    // 3. 简化等待器
    template<typename T>
    concept SchedulerConcept = requires
    {
        { T::add(std::declval<uint32_t>(), std::declval<std::coroutine_handle<> >()) } -> std::same_as<void>;
    };

    template<uint32_t(*getTick)(), SchedulerConcept Scheduler>
    struct TimerAwaitable {
        uint32_t ms_delay; // 无动态时间点存储

        bool await_ready() noexcept { return false; } // 总是挂起

        void await_suspend(const std::coroutine_handle<> h) const
        {
            // 直接注册到调度器
            const auto when = getTick() + ms_delay;
            Scheduler::add(when, h);
        }

        void await_resume() noexcept
        {
        }
    };
}


namespace utils
{
    /**
     * 自定义协程
     * @tparam getTick ms级时钟
     * @tparam MAX_TASKS 最大协程任务数量
     * @tparam FrameSize 静态分配池中，单个协程帧大小
     * @tparam FrameCount 静态分配池中，协程帧数量
     * @note 裸机、FreeRTOS情况下使用需要小心，因为上下文切换并不会保存FPU状态，因此需要避免浮点运算
     * @code{.cpp}
        // 模板实例化
        using coroutine = utils::coroutine<>;


        // 协程任务定义
        coroutine::asynctask blink_coroutine() {
            while (true) {
                std::cout << "LED started" << "\ttime:" << HAL_GetTick() << "\r\n";

                co_await coroutine::awaitable{300}; // 等待300ms
                std::cout << "LED on" << "\ttime:" << HAL_GetTick() << "\r\n";

                co_await coroutine::awaitable{500}; // 等待500ms
                std::cout << "LED off" << "\ttime:" << HAL_GetTick() << "\r\n";

                co_await coroutine::awaitable{600}; // 等待600ms
                std::cout << "LED finished" << "\ttime:" << HAL_GetTick() << "\r\n";
            }
        }

        int main()
        {
            // 协程任务初始化
            blink_coroutine();

            // 事件循环放到主循环里执行
            while (true) {
                coroutine::scheduler_poll<true>();
            }
        }
     * @endcode
     */
    template<uint32_t(*getTick)(), std::size_t MAX_TASKS = 10, size_t FrameSize = 128, size_t FrameCount = 4>
    class coroutine {
        using scheduler = detail_coroutine::Scheduler<getTick, MAX_TASKS>;

    public:
        using awaitable = detail_coroutine::TimerAwaitable<getTick, scheduler>;
        using asynctask = detail_coroutine::AsyncTask<FrameSize, FrameCount>;

        template<bool is_step = false>
        static void scheduler_poll()
        {
            scheduler::template poll<is_step>();
        }
    };
}


#endif //COROUTINE_HPP
